package com.kaige.aiagent.chatMemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.util.ArrayList;
import java.util.List;

public class FileBasedChatMemory implements ChatMemory {
    private final String BASE_DIR;
    // 初始化Kryo序列化工具，用于高效序列化和反序列化对象
    private static final Kryo kryo = new Kryo();

    // 静态代码块配置Kryo实例
    static {
        // 不需要提前注册类即可进行序列化/反序列化操作
        kryo.setRegistrationRequired(false);

        // 使用StdInstantiatorStrategy策略创建对象实例，支持无默认构造函数的类
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedChatMemory(String dir) {
        BASE_DIR = dir;
        File baseDir = new File(dir);
        if (!baseDir.exists()) {
            baseDir.mkdirs();
        }
    }


    @Override
    public void add(String conversationId, List<Message> messages) {
        List<Message> conversationMessages = getOrCreateConversation(conversationId);
        conversationMessages.addAll(messages);
        saveConversation(conversationId, conversationMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> messageList = getOrCreateConversation(conversationId);
        return messageList.stream()
                .skip(Math.max(0, messageList.size() - lastN)) // 跳过前面的消息，只保留最近的lastN条消息
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            file.delete();
        }
    }
    /**
     * 获取指定会话ID对应的聊天消息列表，如果不存在则创建一个新的空列表。
     * 该方法会检查对应会话的文件是否存在，若存在则使用Kryo反序列化读取消息列表；
     * 若不存在，则返回一个空的消息列表。
     *
     * @param conversationId 会话唯一标识
     * @return 返回与会话ID关联的消息列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        List<Message> messages = new ArrayList<>();
        if (file.exists()) {
            try (Input input = new Input(new FileInputStream(file))) {
                messages = kryo.readObject(input, ArrayList.class);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return messages;
    }

    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".kryo");
    }
    private void saveConversation(String conversationId, List<Message> conversationMessages) {
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, conversationMessages);
        }catch (Exception e){
            e.printStackTrace(); // 打印异常堆栈信息，用于调试，显示错误发生的具体位置和异常类型
            System.out.println("发生异常信息：" + e.getMessage()); // 输出简要的异常描述信息，便于快速了解异常原因
        }
    }


}
